from langchain.tools import BaseTool
from langchain.chains import create_sql_query_chain
from langchain_community.utilities import SQLDatabase
from pydantic import BaseModel,Field,ConfigDict
from typing import Type,Any
from model.my_chat_model import ChatModel

#定义工具类输入的参数格式
class ChangePasswordInput(BaseModel):
    email: str = Field(...,description="邮箱")
    old_password: str = Field(...,description="原始密码")
    new_password: str = Field(...,description="新密码")

#定义工具类
class ChangePasswordTool(BaseTool):
    #允许模型接受非标准数据类型
    model_config = ConfigDict(arbitrary_types_allowed=True)
    def __init__(self,**kwargs:Any):
        super().__init__(
            name = "get_changepassword_tool",
            description = "主要用于完成账号密码的修改，必须输入的参数有：邮箱账号、原始密码、新密码",
            **kwargs
        )
    #指定模型的输入参数必须是LoginInput:
    args_schema : Type[BaseModel] = ChangePasswordInput
    #定义工具方法
    def _run(self,email:str,old_password:str,new_password:str):
        #获取大模型
        chat = ChatModel()
        llm = chat.get_line_model()
        #创建Mysql数据库连接
        db = SQLDatabase.from_uri(
            "mysql+pymysql://root:root@localhost:3306/ai",
            include_tables = ["user_info"]
        )
        #创建链
        chain = create_sql_query_chain(llm,db)
        #提问
        question = f"请修改邮箱账号为{email}和原始密码为{old_password}账号的密码，修改后的新密码为{new_password}"
        sql = chain.invoke({"question":question})
        # print(f"sql:{sql}")
        try:
            if "```sql" in sql:
                sql = sql.split("```sql")[1].split("```")[0]
            else:
                sql = sql
            # print(f" .... :{sql}")
            rs=db.run(sql)
            return "密码更新成功"
        except Exception as e:
            return "密码更新失败"
# if __name__ == '__main__':
    # tool=ChangePasswordTool()
    # tool._run(email="84488458@qq.com",old_password="123456789",new_password="abc123456789")